# coding=utf-8
# @author:      ChengJing
# @name:        CNN_model.py
# @datetime:    2021/12/10 15:37
# @software:    PyCharm
# @description:

import torch
import torch.nn as nn


class MyCNN(nn.Module):
    """
    定义简单的卷积神经网络模型进行爆管预警和定位识别
    """
    def __init__(self, in_channel, hidden_channel1, hidden_channel2, out_channel, hidden_feature, out_features):
        """
        Args:
            in_channel: 输入数据的通道数
            hidden_channel1: 隐含层通道数
            hidden_channel2: 隐含层通道数
            out_channel: 卷积输出通道数
            hidden_feature: 隐含层特征
            out_features: 输出特征
        """
        super(MyCNN, self).__init__()
        self.cnn = nn.Sequential(
            nn.Conv2d(in_channels=in_channel, out_channels=hidden_channel1, kernel_size=3, padding=2),
            nn.MaxPool2d(kernel_size=3),
            nn.Conv2d(in_channels=hidden_channel1, out_channels=hidden_channel2, kernel_size=(5, 3), padding=2),
            nn.MaxPool2d(kernel_size=(5, 3)),
            nn.Conv2d(in_channels=hidden_channel2, out_channels=out_channel, kernel_size=(3, 3), padding=2)
        )
        self.adp = nn.AdaptiveMaxPool2d(output_size=1)
        self.classification = nn.Sequential(
            nn.Linear(in_features=out_channel, out_features=hidden_feature),
            nn.ReLU(),
            nn.Linear(in_features=hidden_feature, out_features=out_features),
            nn.Softmax(dim=1)
        )

    def forward(self, x):
        """
        前向传递函数
        """
        x.unsqueeze_(dim=1)  # 将输入数据变成(B, C, S, F) //B-batch size, C-channel size, S-sequence length, F-features
        x = self.cnn(x)
        x = self.adp(x)
        x = self.classification(x.squeeze())
        return x


if __name__ == '__main__':
    model = MyCNN(1, 8, 16, 32, 64, 20)
    from torchsummary import summary
    print(summary(model.cuda(),(60,12)))

